package org.example.platerecognize.service;

import org.apache.commons.lang3.StringUtils;
import org.example.platerecognize.PlateInfoWrapper;
import org.example.platerecognize.PlateVO;
import org.example.platerecognize.core.TorchPlateDetection;
import org.example.platerecognize.core.TorchPlateRecognition;
import org.example.platerecognize.dao.ImageMat;
import org.example.platerecognize.dao.PlateData;
import org.example.platerecognize.utils.CropUtil;
import org.opencv.core.Core;
import org.opencv.core.Mat;
import org.opencv.core.Point;

import java.awt.*;
import java.io.File;
import java.util.HashMap;
import java.util.List;
import java.util.stream.Collectors;

import static org.example.platerecognize.core.TorchPlateDetection.defIouTh;
import static org.example.platerecognize.core.TorchPlateDetection.defScoreTh;
import static org.example.platerecognize.utils.MatUtil.drawPolygon;

public class PlateExtractService {
    private final TorchPlateDetection plateDetection;
    private final TorchPlateRecognition plateRecognition;

    public PlateExtractService(TorchPlateDetection plateDetection, TorchPlateRecognition plateRecognition) {
        this.plateDetection = plateDetection;
        this.plateRecognition = plateRecognition;
    }

    protected List<PlateData> extract(ImageMat image) {
        List<PlateData> plateDatas = plateDetection.inference(image, defScoreTh, defIouTh, new HashMap<>());

        //如果车牌区域占用图片大，会出现为空的
        if (plateDatas.isEmpty()) {
            boolean fast = true;
            if (fast)
                return List.of();
//            System.out.println(">>> plateInfos is Empty ...");
//            printInfo(image.mat(),"原图");

            //将车牌填充为2倍，为了防止由于车牌占用大，导致检测模型识别失败
            Mat tempMat = new Mat();
            int top = image.mat().height() / 2;
            int bot = image.mat().height() / 2;
            int left = image.mat().width() / 2;
            int right = image.mat().width() / 2;
            Core.copyMakeBorder(image.mat(), tempMat, top, bot, left, right, Core.BORDER_CONSTANT);
//            printInfo(tempMat,"填充后的图");

            plateDatas = plateDetection.inference(ImageMat.fromCVMat(tempMat), defScoreTh, defIouTh, new HashMap<>());
            for (PlateData plateData : plateDatas) {
                //将边框扩大，因为车牌被填充为2倍了
                plateData.box = expandPlateBox(image, plateData.box, left, top);
            }
            tempMat.release();
        }

        //取车牌topK
        int topK = 5;
        if (plateDatas.size() > topK) {
            plateDatas = plateDatas.subList(0, topK);
        }

        //解析车牌信息
        for (PlateData plateData : plateDatas) {
            Mat crop = CropUtil.crop(image.mat(), plateData.box);
            plateData.parseInfo = plateRecognition.inference(ImageMat.fromCVMat(crop), plateData.single);
        }

        //清除为空的数据
        plateDatas.removeIf(e ->
                e.parseInfo == null || e.parseInfo.plateNo.length()<5
        );

        Point center = image.center();
        //有效距离
        int distWidth = image.getWidth()/4;
        int distHeight = image.getHeight()/4;

        //返回数据，过滤掉太靠图片边缘的
        return plateDatas.stream()
                .filter(e -> {
                    Point p = e.box.center();
                    var dx = p.x - center.x;
                    var dy = p.y - center.y;
                    return (Math.abs(dx)<=  distWidth && Math.abs(dy)<= distHeight);
                }).toList();
//        return PlateImage.build(image.toBase64AndNoReleaseMat(), plateDatas);
    }

    private static PlateData.PlateBox expandPlateBox(ImageMat image, PlateData.PlateBox plateBox, int left, int top) {
        plateBox = plateBox.clone();

        plateBox = plateBox.move(left, 0, top, 0);
        plateBox.leftTop.x = Math.max(0, plateBox.leftTop.x);
        plateBox.leftTop.y = Math.max(0, plateBox.leftTop.y);
        plateBox.rightTop.x = Math.min(image.getWidth(), plateBox.rightTop.x);
        plateBox.rightTop.y = Math.max(0, plateBox.rightTop.y);
        plateBox.rightBottom.x = Math.min(image.getWidth(), plateBox.rightBottom.x);
        plateBox.rightBottom.y = Math.min(image.getHeight(), plateBox.rightBottom.y);
        plateBox.leftBottom.x = Math.max(0, plateBox.leftBottom.x);
        plateBox.leftBottom.y = Math.min(image.getHeight(), plateBox.leftBottom.y);

        return plateBox;
    }

    private void drawBorder(Mat mat, List<PlateData> plateDatas) {
        for (PlateData plateInfo : plateDatas) {
            var points = plateInfo.box.points();
            drawPolygon(mat,points,Color.RED,2);
        }
    }


    /**
     * 车牌识别
     * 对外暴露的方法
     *
     * @param imageMat
     * @return
     */
    public List<PlateVO> recognition(ImageMat imageMat) {

        List<PlateData> plateDatas = this.extract(imageMat);

        drawBorder(imageMat.mat(), plateDatas);

        String base64 = imageMat.toBase64AndNoReleaseMat();

        return PlateInfoWrapper.wrapPlateResult(plateDatas, base64);
    }

    /**
     * 车牌识别
     * 对外暴露的方法
     *
     * @param file
     * @return
     */
    public List<PlateVO> recognition(File file) {
        ImageMat imageMat = ImageMat.fromImage(file.getPath());
        return recognition(imageMat);
    }

    /**
     * 车牌识别
     * 对外暴露的方法
     *
     * @param path
     * @return
     */
    public List<PlateVO> recognition(String path) {
        ImageMat imageMat = ImageMat.fromImage(path);
        return recognition(imageMat);
    }

}
